import numpy as np
import torch
from torch.utils import data
from d2l import torch as d2l
from data_generate import generate_data

features,labels = generate_data()

def loadArray(data_arrays,batch_size,shuffle=True):
    data_set = data.TensorDataset(*data_arrays)
    return data.DataLoader(data_set,batch_size,shuffle=shuffle)

batch_size = 10
data_iter = loadArray((features,labels),batch_size)

batch = next(iter(data_iter))

from torch import nn
net = nn.Sequential(nn.Linear(2,1))
net[0].weight.data.normal_(0,0.01)
net[0].bias.data.fill_(0)
epochs = 3
alpha = 0.03
loss = nn.MSELoss()
trainer = torch.optim.SGD(net.parameters(),lr=alpha)

def train(epochs,features,labels):
        for epoch in range(epochs):
            for X,y in data_iter:
                l = loss(net(X),y)
                trainer.zero_grad()
                l.backward()
                trainer.step()
            l = loss(net(features),labels)
            print('epoch:',epoch,'loss:',l)

train(epochs,features,labels)